import itertools


# 计算支持度
def calculate_support(dataset, itemset):
    count = 0
    for transaction in dataset:
        if itemset.issubset(transaction):
            count += 1
    return count / len(dataset)


# 生成频繁项集
def apriori(dataset, min_support):
    # 1. 获取所有单项集
    itemsets = [set([item]) for transaction in dataset for item in transaction]
    itemsets = list(set(itemsets))  # 去重

    # 2. 筛选频繁项集
    frequent_itemsets = []
    while itemsets:
        candidate_itemsets = []
        for itemset in itemsets:
            support = calculate_support(dataset, itemset)
            if support >= min_support:
                frequent_itemsets.append((itemset, support))
                candidate_itemsets.append(itemset)

        # 更新为频繁项集的候选项集
        itemsets = []
        for i in range(len(candidate_itemsets)):
            for j in range(i + 1, len(candidate_itemsets)):
                merged_itemset = candidate_itemsets[i] | candidate_itemsets[j]
                if len(merged_itemset) == len(candidate_itemsets[i]) + 1:
                    itemsets.append(merged_itemset)

    return frequent_itemsets


# 示例数据集
dataset = [
    {'milk', 'bread', 'butter'},
    {'milk', 'bread'},
    {'bread', 'butter'},
    {'milk', 'bread', 'butter', 'cheese'},
    {'milk', 'cheese'}
]

# 最小支持度
min_support = 0.6

frequent_itemsets = apriori(dataset, min_support)
for itemset, support in frequent_itemsets:
    print(f"Itemset: {itemset}, Support: {support}")
